import os

import numpy as np
from PIL import Image

# data_path = r'./data/202303/dataset'
data_path = r'./data/20230827'


def gen_dataset():
    def generate_sub_images(ori_path, shade_path, sub_image_size, step_size):
        # 获取文件夹中的所有图像文件
        image_files = [f for f in os.listdir(ori_path) if f.endswith('.jpg')]
        sub_images_ori = []
        sub_images_shade = []
        position = []

        for file_name in image_files:
            img_ori = Image.open(os.path.join(ori_path, file_name))
            img_shade = Image.open(os.path.join(shade_path, file_name))
            width = min(img_ori.size[0], img_shade.size[0])
            height = min(img_ori.size[1], img_shade.size[1])

            # 子图分割
            for y in range(0, height, step_size):
                for x in range(0, width, step_size):
                    # 计算子图的区域
                    region = (x, y, x + sub_image_size, y + sub_image_size)

                    if region[2] > width or region[3] > height:
                        continue

                    # 裁剪子图
                    position.append((x, y))
                    sub_image = img_ori.crop(region)
                    sub_images_ori.append(np.array(sub_image).flatten())
                    sub_image = img_shade.crop(region)
                    sub_images_shade.append(np.array(sub_image).flatten())

            img_ori.close()
            img_shade.close()

        sub_images_ori = np.vstack(sub_images_ori)
        sub_images_shade = np.vstack(sub_images_shade)
        pos_x = np.array([x[0] for x in position])
        pos_y = np.array([x[1] for x in position])

        # sub_images_ori 和 sub_image_shade 的前两列为位置信息x和y，后续为flatten后的子图序列
        sub_images_ori = np.concatenate((pos_x.reshape(-1, 1), pos_y.reshape(-1, 1), sub_images_ori), axis=1)
        sub_images_shade = np.concatenate((pos_x.reshape(-1, 1), pos_y.reshape(-1, 1), sub_images_shade), axis=1)
        return sub_images_ori, sub_images_shade

    # 调用函数进行处理
    print('Start generate dataset.')
    ori = os.path.join(data_path, 'ori')
    shade = os.path.join(data_path, 'shade')
    sub_image_size = 5
    step_size = 3  # 替换为步长
    return generate_sub_images(ori, shade, sub_image_size, step_size)


def process_shade():
    import shutil
    shade = os.path.join(data_path, 'shade')
    outpath = os.path.join(data_path, 'shade2')
    for img in os.listdir(shade):
        img2 = img.replace('-', '')
        img2 = img2.strip('带孔')
        shutil.copy(os.path.join(shade, img), os.path.join(outpath, img2))


def cross_data():
    import shutil
    ori_lst = sorted(os.listdir(os.path.join(data_path, 'ori')))
    shade_lst = sorted(os.listdir(os.path.join(data_path, 'shade2')))

    for i, (img_ori, img_shade) in enumerate(zip(ori_lst, shade_lst)):
        shutil.copy(os.path.join(data_path, 'ori', img_ori), os.path.join(data_path, 'ori_inter', F'{i}.jpg'))
        shutil.copy(os.path.join(data_path, 'shade2', img_shade), os.path.join(data_path, 'shade_inter', F'{i}.jpg'))


if __name__ == '__main__':
    # process_shade()
    # cross_data()
    gen_dataset()
